{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/amrzhd/EEGNet/blob/main/EEGNet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x_5wKiVyBTzL"
      },
      "source": [
        "# Motor Imagery Task Classification\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-3Lg9veBge0"
      },
      "source": [
        "#Extracting Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W4JrI64nBl52"
      },
      "source": [
        "##Downloading BCI Competition IV 2a Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_6jNZi8rBNjN",
        "outputId": "88ac3846-5abc-411f-93f2-f4e96daeed4e"
      },
      "outputs": [],
      "source": [
        "!wget https://www.bbci.de/competition/download/competition_iv/BCICIV_2a_gdf.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5DbydQswB27q"
      },
      "outputs": [],
      "source": [
        "!mkdir -p data/cleaned_data/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "50hgj12yB54U"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!unzip BCICIV_2a_gdf.zip -d data/raw_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y_mkff9aB-9w"
      },
      "source": [
        "#Installing Packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mqHVmwYyCBsB"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install mne"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ueJFbf8KCCPw"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install torch-summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crxgIV9ICD2T"
      },
      "source": [
        "#Libraries Used"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CEp96MytCGei"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import mne\n",
        "import math\n",
        "import copy\n",
        "import gdown\n",
        "import random\n",
        "import scipy.io\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import seaborn as sn\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Torch\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torchsummary import summary\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import DataLoader, TensorDataset, random_split\n",
        "\n",
        "# Scikit-Learn\n",
        "from sklearn.feature_selection import mutual_info_classif\n",
        "from sklearn.metrics import confusion_matrix, accuracy_score\n",
        "from sklearn.model_selection import train_test_split"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nn8IHP2_CQbH"
      },
      "source": [
        "# Building Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JFtUH89-G48b"
      },
      "outputs": [],
      "source": [
        "!mkdir -p data/cleaned_data/first_session"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6Bmm_acGc4g"
      },
      "source": [
        "##First Session"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "X6ViqHGaCUWY",
        "outputId": "39111580-ee4e-4543-f906-f9de7e97d6c7"
      },
      "outputs": [],
      "source": [
        "raw_data_folder = 'data/raw_data/'\n",
        "cleaned_data_folder = 'data/cleaned_data/first_session/'\n",
        "files = os.listdir(raw_data_folder)\n",
        "\n",
        "# Selecting files with suffix 'T.gdf'\n",
        "filtered_files = [file for file in files if file.endswith('T.gdf')]\n",
        "\n",
        "raw_list = []\n",
        "\n",
        "# Iterating through filtered files \n",
        "'''\n",
        "    使用MNE处理EEG数据\n",
        "    1.去眼动(EOG)噪声\n",
        "    2.滤波\n",
        "    3.去基线漂移\n",
        "'''\n",
        "for file in filtered_files:\n",
        "    file_path = os.path.join(raw_data_folder, file)\n",
        "\n",
        "    # Reading raw data\n",
        "    raw = mne.io.read_raw_gdf(file_path, eog=['EOG-left', 'EOG-central', 'EOG-right'], preload=True)\n",
        "    # Droping EOG channels ['EOG-left', 'EOG-central', 'EOG-right']\n",
        "    #### 1.implement your code here \n",
        "\n",
        "\n",
        "    # High Pass Filtering 4-40 Hz\n",
        "    #### 2.implement your code here  \n",
        "\n",
        "    # Notch filter for Removal of Line Voltage\n",
        "\n",
        "    #### 3.implement your code here  \n",
        "\n",
        "    # Saving the modified raw data to a file with .fif suffix\n",
        "    new_file_path = os.path.join(cleaned_data_folder, file[:-4] + '.fif')\n",
        "    raw.save(new_file_path, overwrite=True)\n",
        "    # Appending data to the list\n",
        "    raw_list.append(raw)\n",
        "\n",
        "final_raw = mne.concatenate_raws(raw_list)\n",
        "new_file_path = os.path.join(cleaned_data_folder, 'First_Session_Subjects.fif')\n",
        "final_raw.save(new_file_path, overwrite=True)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wNqq3V3rCiPn"
      },
      "source": [
        "**List of the events**  \n",
        "'1023': 1 Rejected trial  \n",
        "'1072': 2 Eye movements  \n",
        "'276': 3 Idling EEG (eyes open)  \n",
        "'277': 4 Idling EEG (eyes closed)  \n",
        "'32766': 5 Start of a new run  \n",
        "'768': 6 Start of a trial  \n",
        "'769': 7 Cue onset **Left** (class 1) : 0  \n",
        "'770': 8 Cue onset **Right** (class 2) : 1  \n",
        "'771': 9 Cue onset **Foot** (class 3) : 2  \n",
        "'772': 10 Cue onset **Tongue** (class 4): 3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-9RZjQbWCoc0",
        "outputId": "fbbd728f-9f7f-457e-d4c9-eaefe67bafa1"
      },
      "outputs": [],
      "source": [
        "events = mne.events_from_annotations(final_raw)\n",
        "events[1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PJAI_olujhxI"
      },
      "source": [
        "**Time choice:**  \n",
        "[0.s, 4s] Post Cue on set:  [2s, 6s]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OaAuA1mQCqcP",
        "outputId": "4d33fb54-48f8-4eb9-c516-1b97747e39d6"
      },
      "outputs": [],
      "source": [
        "epochs = mne.Epochs(final_raw, events[0], event_id=[7, 8, 9, 10], tmin=0, tmax=4, reject=None, baseline=None, preload=True)\n",
        "# 读取epochs中的数据和标签\n",
        "### implement your code here \n",
        "first_session_data = \n",
        "first_session_labels = "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rVNRjQbfCs3-",
        "outputId": "041d64ef-2b44-49e8-d554-90bc0992bc36"
      },
      "outputs": [],
      "source": [
        "print(\"First_session_dataset shape:\",first_session_data.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nA0Hr5BMPE5e"
      },
      "source": [
        "###Structuring Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D5dF6ZydC2PP",
        "outputId": "67fafc85-ff3d-48f3-f08a-3abc663df0b3"
      },
      "outputs": [],
      "source": [
        "# Choosing Device\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Loss Function classification tasks\n",
        "#### 5.implement your code here  \n",
        "\n",
        "# Normalizing Labels to [0, 1, 2, 3]\n",
        "y_train = first_session_labels - np.min(first_session_labels)\n",
        "\n",
        "# Normalizing Input features: z-score(mean=0, std=1)\n",
        "X_first_session = (first_session_data - np.mean(first_session_data)) / np.std(first_session_data)\n",
        "\n",
        "X = X_first_session\n",
        "y = y_train\n",
        "\n",
        "# Spliting  Data: 90% for Train and 10% for Test\n",
        "#### 6.implement your code here  #### \n",
        "\n",
        "\n",
        "# Converting to Tensor\n",
        "X_train = torch.Tensor(X_train).unsqueeze(1).to(device)\n",
        "X_test = torch.Tensor(X_test).unsqueeze(1).to(device)\n",
        "y_train = torch.LongTensor(y_train).to(device)\n",
        "y_test = torch.LongTensor(y_test).to(device)\n",
        "\n",
        "# Creating Tensor Dataset\n",
        "#### 7.implement your code here  \n",
        "\n",
        "# Printing the sizes\n",
        "print(\"Size of X_train:\", X_train.size())\n",
        "print(\"Size of X_test:\", X_test.size())\n",
        "print(\"Size of y_train:\", y_train.size())\n",
        "print(\"Size of y_test:\", y_test.size())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xyXDZiZ4E4a_"
      },
      "source": [
        "#Training Class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F3nUwuQHE8Hi"
      },
      "outputs": [],
      "source": [
        "class TrainModel():\n",
        "    def __init__(self,):\n",
        "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "    def train_model(self, model, train_dataset, learning_rate=0.001, batch_size=64, epochs=500):\n",
        "\n",
        "        model = model.to(self.device)\n",
        "        criterion = nn.CrossEntropyLoss()\n",
        "        optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "        highest_train_accuracy = 0.0\n",
        " \n",
        "        highest_train_accuracy = 0.0\n",
        "\n",
        "        for epoch in range(epochs):\n",
        "            model.train()\n",
        "            running_loss = 0.0\n",
        "            correct = 0\n",
        "            total = 0\n",
        "            for inputs, labels in train_loader:\n",
        "                inputs = inputs.to(self.device)\n",
        "                labels = labels.to(self.device)\n",
        "\n",
        "                #### 8.implement your code here model/criterion/optimizer/data_loader\n",
        "\n",
        "                running_loss += loss.item() * inputs.size(0)\n",
        "                _, predicted = torch.max(outputs, 1)\n",
        "                total += labels.size(0)\n",
        "                correct += (predicted == labels).sum().item()\n",
        "\n",
        "            epoch_loss = running_loss / len(train_loader.dataset)\n",
        "            epoch_accuracy = correct / total\n",
        "            if epoch_accuracy > highest_train_accuracy:\n",
        "                highest_train_accuracy = epoch_accuracy\n",
        "            print(f\"Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {(epoch_accuracy*100):.2f}%\")\n",
        "\n",
        "        average_loss = running_loss / len(train_loader.dataset)\n",
        "        print(\"Average Loss:\", average_loss)\n",
        "        print(\"Highest Train Accuracy:\", highest_train_accuracy)\n",
        "\n",
        "        # Saving model\n",
        "        torch.save(model.state_dict(), 'eegnet_model.pth')\n",
        "        return model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ObMKFBTxFLBJ"
      },
      "source": [
        "#Evaluating Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GxVYOQ_YFM4b"
      },
      "outputs": [],
      "source": [
        "class EvalModel():\n",
        "    def __init__(self, model):\n",
        "        self.model = model.to(device)\n",
        "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "    def test_model(self, test_dataset):\n",
        "        self.model.eval()\n",
        "        correct = 0\n",
        "        total = 0\n",
        "        test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for inputs, labels in test_loader:\n",
        "                inputs = inputs.to(self.device)\n",
        "                labels = labels.to(self.device)\n",
        "                outputs = self.model(inputs)\n",
        "                _, predicted = torch.max(outputs.data, 1)\n",
        "                #### 9.implement your code here 计算准确率\n",
        "   \n",
        "\n",
        "        accuracy = (correct / total) * 100\n",
        "        print(\"/------------------------------/\")\n",
        "        print(f\"Test Accuracy: {accuracy:.2f}%\")\n",
        "        print(\"/------------------------------/\")\n",
        "        return accuracy\n",
        "\n",
        "    def plot_confusion_matrix(self, test_dataset, classes):\n",
        "        self.model.eval()\n",
        "        y_pred = []\n",
        "        y_true = []\n",
        "        test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            for inputs, labels in test_loader:\n",
        "                inputs = inputs.to(self.device)\n",
        "                labels = labels.to(self.device)\n",
        "                outputs = self.model(inputs)\n",
        "                _, predicted = torch.max(outputs.data, 1)\n",
        "                y_pred.append(predicted.item())\n",
        "                y_true.append(labels.item())\n",
        "\n",
        "        cf_matrix = confusion_matrix(y_true, y_pred)\n",
        "        cf_matrix = cf_matrix.astype('float') / cf_matrix.sum(axis=1)[:, np.newaxis]\n",
        "\n",
        "        df_cm = pd.DataFrame(cf_matrix, index=classes, columns=classes)\n",
        "\n",
        "        plt.figure(figsize=(10, 7))\n",
        "        sn.heatmap(df_cm, annot=True, cmap='Blues', fmt='.2f')\n",
        "        plt.xlabel('Predicted labels')\n",
        "        plt.ylabel('True labels')\n",
        "        plt.title('Confusion Matrix')\n",
        "        plt.savefig('confusion_matrix_model.png')\n",
        "        plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vR2z57zqC5iS"
      },
      "source": [
        "#EEGNet Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SmFhrHzSC8BE"
      },
      "outputs": [],
      "source": [
        "class EEGNetModel(nn.Module): # EEGNET-8,2\n",
        "    def __init__(self, chans=22, classes=4, time_points=1001, temp_kernel=32,\n",
        "                 f1=16, f2=32, d=2, pk1=8, pk2=16, dropout_rate=0.5, max_norm1=1, max_norm2=0.25):\n",
        "        super(EEGNetModel, self).__init__()\n",
        "        # Calculating FC input features\n",
        "        linear_size = (time_points//(pk1*pk2))*f2\n",
        "\n",
        "        # Temporal Filters\n",
        "        self.block1 = nn.Sequential(\n",
        "            nn.Conv2d(1, f1, (1, temp_kernel), padding='same', bias=False),\n",
        "            nn.BatchNorm2d(f1),\n",
        "        )\n",
        "        # Spatial Filters\n",
        "        self.block2 = nn.Sequential(\n",
        "            nn.Conv2d(f1, d * f1, (chans, 1), groups=f1, bias=False), # Depthwise Conv\n",
        "            nn.BatchNorm2d(d * f1),\n",
        "            nn.ELU(),\n",
        "            nn.AvgPool2d((1, pk1)),\n",
        "            nn.Dropout(dropout_rate)\n",
        "        )\n",
        "        self.block3 = nn.Sequential(\n",
        "            nn.Conv2d(d * f1, f2, (1, 16),  groups=f2, bias=False, padding='same'), # Separable Conv\n",
        "            nn.Conv2d(f2, f2, kernel_size=1, bias=False), # Pointwise Conv\n",
        "            nn.BatchNorm2d(f2),\n",
        "            nn.ELU(),\n",
        "            nn.AvgPool2d((1, pk2)),\n",
        "            nn.Dropout(dropout_rate)\n",
        "        )\n",
        "        self.flatten = nn.Flatten()\n",
        "        self.fc = nn.Linear(linear_size, classes)\n",
        "\n",
        "        # Apply max_norm constraint to the depthwise layer in block2\n",
        "        self._apply_max_norm(self.block2[0], max_norm1)\n",
        "\n",
        "        # Apply max_norm constraint to the linear layer\n",
        "        self._apply_max_norm(self.fc, max_norm2)\n",
        "\n",
        "    def _apply_max_norm(self, layer, max_norm):\n",
        "        for name, param in layer.named_parameters():\n",
        "            if 'weight' in name:\n",
        "                param.data = torch.renorm(param.data, p=2, dim=0, maxnorm=max_norm)\n",
        "\n",
        "    def forward(self, x):\n",
        "         #### 10.implement your code here \n",
        "        '''\n",
        "            完成forward操作\n",
        "            提示：注意输入输出的tensor维度关系\n",
        "        '''\n",
        "        x = self.fc(x)\n",
        "        return x\n",
        "    \n",
        "    ####附加题\n",
        "    #### 11.请参考上述代码，设计一个基于1D Conv/Transformer/Mamba/等其他深度学习模型的脑电信号分类模型，并亲自跑出结果\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M49ks2b8C_dP"
      },
      "source": [
        "##Model Summery"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7NdOeIRtDAOn",
        "outputId": "35ec6ad5-7e2c-4f99-cae0-e83ffd549797"
      },
      "outputs": [],
      "source": [
        "input_size = (1, 22, 1001)\n",
        "eegnet_model = EEGNetModel().to(device)\n",
        "summary(eegnet_model, input_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "95r4EymoDCIv"
      },
      "source": [
        "##Training Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TJzl6On9DFcp",
        "outputId": "7c7e7ca1-e094-478a-f7df-171e12388792"
      },
      "outputs": [],
      "source": [
        "eegnet_model = EEGNetModel().to(device)\n",
        "\n",
        "# Training Hyperparameters\n",
        "EPOCHS = 500\n",
        "BATCH_SIZE = 64\n",
        "LEARNING_RATE = 0.001\n",
        "trainer = TrainModel()\n",
        "trained_eegnet_model = trainer.train_model(eegnet_model, train_dataset, learning_rate=LEARNING_RATE,\n",
        "                                   batch_size=BATCH_SIZE, epochs=EPOCHS)\n",
        "torch.save(trained_eegnet_model.state_dict(), 'eegnet_model.pth')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3fKnNHxDKwn"
      },
      "source": [
        "##Evaluating Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 695
        },
        "id": "aHsstvZSDLb1",
        "outputId": "62367285-83e7-4ed4-db94-376f791fcfec"
      },
      "outputs": [],
      "source": [
        "classes_list = ['Left', 'Right', 'Foot', 'Tongue']\n",
        "eval_model = EvalModel(trained_eegnet_model)\n",
        "test_accuracy = eval_model.test_model(test_dataset)\n",
        "eval_model.plot_confusion_matrix(test_dataset, classes_list)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "authorship_tag": "ABX9TyOuzENxctn49HNjNmnJOE11",
      "gpuType": "T4",
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "EEG2Image",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
